# -*- coding: utf-8 -*-
# @Time    : 2024/12/20 15:46
# @Author  : yujiahao
# @File    : olama_ask.py
# @description:
import json
import logging
import requests
from fastapi import FastAPI, HTTPException, Request
from ollama import chat

app = FastAPI()

# 配置日志
logging.basicConfig(level=logging.INFO)

# 存储已处理的消息ID
processed_messages = set()


class QuestionRequest:
    def __init__(self, question: str):
        self.question = question


def get_access_token(app_id: str, app_secret: str) -> str:
    url = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal/"
    headers = {
        "Content-Type": "application/json"
    }
    payload = {
        "app_id": app_id,
        "app_secret": app_secret
    }
    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 200:
        data = response.json()
        return data.get("tenant_access_token")
    else:
        logging.error(f"Failed to get access token: {response.text}")
        return None


async def send_message_to_feishu(open_id: str, content: str, access_token: str):
    url = "https://open.feishu.cn/open-apis/message/v4/send/"
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    }
    data = {
        "open_id": open_id,
        "msg_type": "text",
        "content": {
            "text": content
        }
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code != 200:
        logging.error(f"Failed to send message to Feishu: {response.text}")


@app.post("/feishu/askai")
async def handle_request(request: Request):
    try:
        body_bytes = await request.body()
        body = json.loads(body_bytes)
        logging.info(f"Received request body: {body}")

        if "challenge" in body:
            return {"challenge": body["challenge"]}

        if "event" in body:
            event_data = body["event"]
            message_id = event_data["message"]["message_id"]

            # 检查消息是否已处理
            if message_id in processed_messages:
                logging.info(f"Message {message_id} has already been processed.")
                return {"status": "Message already processed"}

            # 标记消息为已处理
            processed_messages.add(message_id)

            if "message" in event_data and "content" in event_data["message"]:
                message_content = json.loads(event_data["message"]["content"])
                question_request = QuestionRequest(question=message_content["text"])

                # 调用 ollama 的 chat 函数
                response = chat(
                    # model="llama3.2",
                    model="deepseek-r1:8b",
                    messages=[{"role": "user", "content": question_request.question}]
                )

                response_content = response["message"]["content"]
                logging.info(f"Response content: {response_content}")

                # 获取 open_id 并发送消息回飞书
                open_id = event_data["sender"]["sender_id"]["open_id"]
                app_id = "XXX"
                app_secret = "XXX"
                access_token = get_access_token(app_id, app_secret)

                if access_token:
                    await send_message_to_feishu(open_id, response_content, access_token)

                return {"answer": response_content}

        return {"error": "Invalid event format"}
    except Exception as e:
        logging.error(f"Error processing request: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")


@app.get("/health")
async def get_status():
    return {"status": "Service is running"}
